from osgeo import gdal,ogr,osr,gdalconst
from xml.dom.minidom import parse
from datetime import datetime
import zipfile
import shutil
import json
import sys
import os
import re


def zipfile_2_jp2files(zip_path, bands, save_dir):
    band_names = {
        "1":"T.*_B01\.jp2",
        "2":"T.*_B02\.jp2",
        "3":"T.*_B03\.jp2",
        "4":"T.*_B04\.jp2",
        "5":"T.*_B05\.jp2",
        "6":"T.*_B06\.jp2",
        "7":"T.*_B07\.jp2",
        "8A":"T.*_B8A\.jp2",
        "8":"T.*_B08\.jp2",
        "9":"T.*_B09\.jp2",
        "10":"T.*_B10\.jp2",
        "11":"T.*_B11\.jp2",
        "12":"T.*_B12\.jp2"}
    pathres = list(map(lambda x:".*/IMG_DATA/"+band_names[f"{x}"],bands))
    paths = []
    zipFile = zipfile.ZipFile(zip_path, "r")
    for pathre in pathres:
        for info in zipFile.infolist():
            if re.match(pathre,info.filename):
                path = zipFile.extract(info.filename, save_dir)
                paths.append(path)
    
    zipFile.close()
    return paths


def read_xmlfile(xml_path,select_node):
    domTree = parse(xml_path)
    collection = domTree.documentElement
    select_element = [collection]
    for node_dict in select_node:
        tag_name = node_dict["tag_name"]
        if node_dict["attribute"] is None:
            select_element = [
                element_sub
                for element in select_element
                for element_sub in element.getElementsByTagName(tag_name)
            ]
        else:
            attribute_name = node_dict["attribute"]["name"]
            attribute_value = node_dict["attribute"]["value"]
            select_element = [
                element_sub
                for element in select_element
                for element_sub in element.getElementsByTagName(tag_name)
                if element_sub.getAttribute(attribute_name) == attribute_value
            ]
    datas = [
        element.firstChild.data
        for element in select_element
    ]
    return datas


def read_cloud(zip_path):
    select_node = [
        {
            "tag_name": "n1:Quality_Indicators_Info",
            "attribute": None
        },
        {
            "tag_name": "Cloud_Coverage_Assessment",
            "attribute": None
        }
    ]
    xml_namere = "^S2.*\.SAFE/MTD_MSIL1C\.xml$"
    zip_file = zipfile.ZipFile(zip_path, "r")
    cloud = "00"
    for file_info in zip_file.infolist():
        if re.match(xml_namere, file_info.filename):
            xml_path = zip_file.open(file_info.filename, "r")
            datas = read_xmlfile(xml_path, select_node)
            cloud = round(float(datas[0]))
            cloud = f"0{cloud}" if len(f"{cloud}") == 1 else f"{cloud}"
    return float(cloud)


def projection_2_spatil(tif_path):
    # 打开栅格数据集
    dataset = gdal.Open(tif_path)
    # 获取WKT格式的空间参考信息
    wkt = dataset.GetProjection()
    # 创建空间参考对象
    spatial_ref = osr.SpatialReference(wkt=wkt)
    # 检查空间参考是否存在
    if wkt == '':
        raise "无空间参考信息"
    else:
        # 检查是否为地理坐标系
        if spatial_ref.IsGeographic():
            print("使用的是地理坐标系")
            # 获取EPSG代码
            epsg_code = spatial_ref.GetAttrValue('AUTHORITY', 1)
            # 检查是否是WGS84
            if epsg_code == '4326':
                print("该栅格使用的是WGS84坐标系。")
                # vrt_path = re.sub("\..*$",".vrt",tif_path)
                mem_path = '/vsimem/output.tif'
                gdal.Warp(mem_path, dataset, dstSRS='EPSG:4326')
                return mem_path
            else:
                # vrt_path = re.sub("\..*$",".vrt",tif_path)
                mem_path = '/vsimem/output.tif'
                # 使用gdal.Warp()进行重新投影
                gdal.Warp(mem_path, dataset, dstSRS='EPSG:4326')
                return mem_path
        # 检查是否为投影坐标系
        elif spatial_ref.IsProjected():
            print("使用的是投影坐标系")
            # 使用gdal.Warp()进行重新投影
            # vrt_path = re.sub("\..*$",".vrt",tif_path)
            mem_path = '/vsimem/output.tif'
            gdal.Warp(mem_path, dataset, dstSRS='EPSG:4326')
            return mem_path


def jp2files_2_tiffile(jp2_paths,tif_path):
    mem_path = '/vsimem/output.tif'
    options = gdal.BuildVRTOptions(
        resolution="highest",
        resampleAlg=gdalconst.GRA_NearestNeighbour,
        separate=True)
    gdal.BuildVRT(mem_path,jp2_paths,options=options)
    creationOptions = [
        "BIGTIFF=YES",
        "TILED=YES",
        "BLOCKXSIZE=1024",
        # "TFW=YES",
        "NUM_THREADS=ALL_CPUS"]
    options = gdal.TranslateOptions(
        format="GTiff",
        xRes=10,
        yRes=10,
        resampleAlg=gdalconst.GRA_NearestNeighbour,
        creationOptions=creationOptions)
    gdal.Translate(tif_path,mem_path,options=options)
    return


def rasterfile_2_geometry(tif_path):
    # 打开栅格数据集
    raster_path = projection_2_spatil(tif_path)
    dataset = gdal.Open(raster_path,gdalconst.GA_Update)
    wkt_projection = dataset.GetProjection()
    band = dataset.GetRasterBand(1)
    nodata = band.GetNoDataValue()
    array = band.ReadAsArray()
    mask_band = band.GetMaskBand()
    mask_array = mask_band.ReadAsArray()
    band.WriteArray(mask_array)
    band = dataset.GetRasterBand(1)
    srs = osr.SpatialReference()
    srs.ImportFromWkt(wkt_projection)
    # 创建输出矢量数据源
    driver = ogr.GetDriverByName('Memory') # Memory
    mem_ds = driver.CreateDataSource('memory') # memory
    layer = mem_ds.CreateLayer('polygon_layer', srs, ogr.wkbPolygon)
    # 添加字段
    fieldDef = ogr.FieldDefn('bandvalue', ogr.OFTInteger)
    layer.CreateField(fieldDef)
    options = []
    # options.append('8CONNECTED=4')  # 使用 8 连通区域生长算法
    gdal.Polygonize(
        band, 
        mask_band, 
        layer, 
        -1, 
        options=options, 
        callback=gdal.TermProgress_nocb
    )
    fearture = layer.GetFeature(0) 
    geometry = fearture.geometry()
    geometry_wkt = geometry.ExportToWkt()
    mem_ds.Destroy()
    del dataset
    return geometry_wkt


def zipfile_2_rasterfile(zip_path,save_dir,bands):
    zip_name = os.path.basename(zip_path)
    yearmonth = zip_name.split('_')[2][:6]
    month_dir = os.path.join(save_dir, yearmonth)
    if not os.path.exists(month_dir):os.makedirs(month_dir)
    tif_name = zip_name.replace('.zip','.tif')
    tif_path = os.path.join(month_dir, tif_name)
    print(tif_path)
    if os.path.exists(tif_path): return tif_path
    jp2_paths = zipfile_2_jp2files(zip_path,bands,month_dir)
    jp2files_2_tiffile(jp2_paths,tif_path)
    name = os.path.basename(zip_path).replace(".zip",".SAFE")
    unzip_dir = os.path.join(month_dir,name)
    shutil.rmtree(unzip_dir)
    return tif_path


def main(task_id,file_path,save_dir):
    date_dir = r"/data/RS_Data/Image_Data/Sentinel2/1C"
    bands = ['2','3','4','5','6','7',"8A",'8','11','12']
    if os.path.splitext(file_path)[-1] == '.zip':
        zip_name = os.path.basename(file_path)
        infos = zip_name.split("_")
        tile = infos[5][1:]
        product_date = infos[2]
        rst_path = zipfile_2_rasterfile(file_path,date_dir,bands) 
        geometry = rasterfile_2_geometry(rst_path)
        cloud = read_cloud(file_path)
        data = {
            "id": task_id,
            "tile": tile,  # 假设其他属性字段在结果中的索引为1
            "product_date": product_date,
            "cpjbdm": "1",
            "cpjdmc": "L1C",
            "imgtype_code": "S2",
            "imgtype_name": "Sentinel-2",
            "cloud": cloud,
            "file_path": file_path,
            "tif_path": rst_path,
            "geometry": geometry
        } 
    else:
        geometry = rasterfile_2_geometry(file_path)
        print(geometry)
        data = {
            "id":task_id,
            "file_path": file_path,
            "geometry": geometry
        }
    json_path = os.path.join(save_dir,"result.json")
    with open(json_path, "w") as json_file:
        json.dump(data, json_file, indent=4)
    return


if __name__ == "__main__":
    task_id = r"test11"
    # file_path = r"/data/FengYang/data/test_data/S2A_MSIL1C_20240302T030631_N0510_R075_T50SME_20240302T051154.zip"
    file_path = r"/data/fengyy/dateset/temp/202403/131123_武强县20230726.tif"
    save_dir = r"/data/fengyy/dateset/temp/202403"
    # main(sys.argv[1],sys.argv[2],sys.argv[3])
    main(task_id,file_path,save_dir)